{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Collision Avoidance - Build TensorRT model for live demo\n",
    "\n",
    "In this notebook we'll use the model we trained to detect whether the robot is ``free`` or ``blocked`` to enable a collision avoidance behavior on the robot.  \n",
    "\n",
    "## Load the trained model\n",
    "\n",
    "We'll assumed that you've already downloaded the ``best_model.pth`` to your workstation as instructed in the training notebook.  Now, you should upload this model into this notebook's\n",
    "directory by using the Jupyter Lab upload tool.  Once that's finished there should be a file named ``best_model.pth`` in this notebook's directory.  \n",
    "\n",
    "> Please make sure the file has uploaded fully before calling the next cell\n",
    "\n",
    "Execute the code below to initialize the PyTorch model.  This should look very familiar from the training notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "model = torchvision.models.resnet18(pretrained=False)\n",
    "model.fc = torch.nn.Linear(512, 2)\n",
    "model = model.cuda().eval().half()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, load the trained weights from the ``best_model_resnet18.pth`` file that you uploaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load('best_model_resnet18.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, the model weights are located on the CPU memory execute the code below to transfer to the GPU device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> If your setup does not have `torch2trt` installed, you need to first install `torch2trt` by executing the following in the console.\n",
    "```bash\n",
    "cd $HOME\n",
    "git clone https://github.com/NVIDIA-AI-IOT/torch2trt\n",
    "cd torch2trt\n",
    "sudo python3 setup.py install\n",
    "```\n",
    "\n",
    "Convert and optimize the model using torch2trt for faster inference with TensorRT. Please see the torch2trt readme for more details.\n",
    "\n",
    "> This optimization process can take a couple minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch2trt import torch2trt\n",
    "\n",
    "data = torch.zeros((1, 3, 224, 224)).cuda().half()\n",
    "\n",
    "model_trt = torch2trt(model, [data], fp16_mode=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the optimized model using the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model_trt.state_dict(), 'best_model_trt.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next\n",
    "\n",
    "Open live_demo_resnet18_build_trt.ipynb to move JetBot with the TensorRT optimized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
